#!usr/bin/env python
#Runs everything in the project. Input is a project dir,database name and then the dirs with the tree and the derived multialignments.
from Bio import SeqIO
from Bio.Seq import Seq
import sys
import os
import re
import shutil
import dendropy
import sqlite3
import numpy
import types

import remove_columns

def check_tree(tree,sql_cursor):
	'''Check if tree in newick format already is in database sql_cursor. Returns true if its in the database and the id for that tree, if its not in the database it returns false and the maximum id in database+1 as a new tree id.'''
	t=(str(tree),)
	sql_cursor.execute('''select tree_id from Tree where newick_format like ?''',t)
	database_tree='not'
	for row in sql_cursor:
		database_tree=row
	if database_tree!='not':
		tree_id=database_tree[0]
		return(True,tree_id)
	else:
		sql_cursor.execute('''select max(tree_id) from Tree''')
		for row in sql_cursor:
			m=row
		if isinstance(m[0],types.NoneType):
			tree_id=1
		else:
			tree_id=m[0]+1	
		return(False,tree_id)

	

sys.stderr.write('Starting\n')
pattern_tree=re.compile('.*.tree')
pattern_fasta=re.compile('.*.msl')
pattern_dot=re.compile('\..*')
sys.argv.pop(0) #Remove first item in sys.argv since it is program name.
if len(sys.argv)<3:
	sys.stderr.write('To few arguments\n')
	sys.exit()
project_dir=sys.argv.pop(0)

#Create dir Results and level below
path_result_dir=project_dir+'Results/'
if not os.path.isdir(path_result_dir):
	os.mkdir(path_result_dir)
path_remove_align_dir=path_result_dir+'Removed_columns_align/'
if not os.path.isdir(path_remove_align_dir):
	os.mkdir(path_remove_align_dir)
path_remove_dist_dir=path_result_dir+'Removed_dist_matrix/'
if not os.path.isdir(path_remove_dist_dir):
	os.mkdir(path_remove_dist_dir)
path_org_dist_dir=path_result_dir+'Org_dist_matrix/'
if not os.path.isdir(path_org_dist_dir):
	os.mkdir(path_org_dist_dir)
path_remove_tree_dir=path_result_dir+'Removed_tree/'
if not os.path.isdir(path_remove_tree_dir):
	os.mkdir(path_remove_tree_dir)
path_org_tree_dir=path_result_dir+'Org_tree/'
if not os.path.isdir(path_org_tree_dir):
	os.mkdir(path_org_tree_dir)

#Create sqlite3 databse if it dosen't exist otherwise ask if program should quit or continue
sql_database_path=sys.argv.pop(0)
if os.path.isfile(sql_database_path):
	sys.stderr.write('Database: '+sql_database_path+' alreday exits.\n Use this database?(Y/N):')
	q=sys.stdin.readline()
	if q=='N\n':
		sys.exit()
	elif q=='Y\n':
		sql_database=sqlite3.connect(sql_database_path)
		sql_cursor=sql_database.cursor()
	else:
		sys.stderr.write('Wrong input\n')
		sys.exit()
else:
	sql_database=sqlite3.connect(sql_database_path)
	sql_cursor=sql_database.cursor()
	sql_cursor.execute('''create table Tree (
	tree_id int primary key,
	newick_format character varying(1000) unique
)''')
	sql_cursor.execute('''create table Distances (
	expr_id int,
	base_tree_id int,
	tree_id int,
	distance float,
	type character varying(20)
)''')
	sql_cursor.execute('''create table Experiments (
	expr_id int primary key,
	description character varying(2000)
)''')
	sql_cursor.execute('''create table Statistics (
	expr_id int primary key,
	min_dist_remove float,
	min_dist_org float,
	max_dist_remove float,
	max_dist_org float,
	median_dist_remove float,
	median_dist_org float,
	mean_dist_remove float,
	mean_dist_org float
)''')
	sql_database.commit()

for dir_path in sys.argv:
	sys.stderr.write('Working on dir: '+dir_path+'\n') 
	try:
		dirList=os.listdir(dir_path) #Make list of files in dir
	except:
		sys.stderr.write('Directory '+dir_path+' doesn\'t exist\n')
		sys.exit()
	if dir_path[len(dir_path)-1]=='/':
		dir_path_name=dir_path.split('/')
		dir_path_name=dir_path_name[len(dir_path_name)-2]+'/'
	else:
		dir_path_name=dir_path.rpartition('/')[2]+'/'
		dir_path=dir_path+'/'
	#Add dir name to description field in database


	sql_cursor.execute('''select max(expr_id) from Experiments''')
	for row in sql_cursor:
		m=row
	if isinstance(m[0],types.NoneType):
		num=1
	else:
		num=m[0]+1
	t=(num,dir_path_name,)
	sql_cursor.execute('insert into Experiments values(?,?)',t)
	expr_id=sql_cursor.lastrowid

	#Create dirs in Results dirs
	path_remove_align_sub_dir=path_remove_align_dir+dir_path_name
	if not os.path.isdir(path_remove_align_sub_dir):
		os.mkdir(path_remove_align_sub_dir)
	path_remove_dist_sub_dir=path_remove_dist_dir+'/'+dir_path_name
	if not os.path.isdir(path_remove_dist_sub_dir):
		os.mkdir(path_remove_dist_sub_dir)
	path_org_dist_sub_dir=path_org_dist_dir+'/'+dir_path_name
	if not os.path.isdir(path_org_dist_sub_dir):
		os.mkdir(path_org_dist_sub_dir)
	path_remove_tree_sub_dir=path_remove_tree_dir+'/'+dir_path_name
	if not os.path.isdir(path_remove_tree_sub_dir):
		os.mkdir(path_remove_tree_sub_dir)
	path_org_tree_sub_dir=path_org_tree_dir+'/'+dir_path_name
	if not os.path.isdir(path_org_tree_sub_dir):
		os.mkdir(path_org_tree_sub_dir)
	
	tree_removed_list=[]
	tree_org_list=[]
	tree_removed_id_list=[]
	tree_org_id_list=[]
	for file_name in dirList: #Loop trough files in dir
   		if pattern_tree.search(file_name) and not pattern_dot.match(file_name): #Check if its the tree
			sys.stderr.write('Working on: '+file_name+'\n')
			tree_base=dendropy.Tree.get_from_path(dir_path+file_name,schema='newick')

			#Insert tree into sql database
			[test,tree_base_id]=check_tree(tree_base,sql_cursor) #Check if tree is already in database and get the id if it is or the maximum id+1 if it isnt
			if not test:
				t=(tree_base_id,str(tree_base),)
				sql_cursor.execute('''insert into Tree values(?,?)''',t)
			tree_base_bool=True #Program will stop later if this is not true
		elif pattern_fasta.search(file_name) and not pattern_dot.match(file_name): #Check if its a fasta file
			sys.stderr.write('Working on: '+file_name+'\n')
			path_new_file=remove_columns.main(dir_path+file_name) #Remove noisy columns

			#Make distance matrix and create trees
			#Removed columns alignment distance matrix
			output_file_remove_dist=path_remove_dist_sub_dir+file_name.rpartition('.')[0]+'.dist'
			os.system('/info/appbio10/bin/fastprot -o '+output_file_remove_dist+' --output-format=phylip '+path_new_file)
			#Org alignment distance matrix
			output_file_org_dist=path_org_dist_sub_dir+file_name.rpartition('.')[0]+'.dist'
			os.system('/info/appbio10/bin/fastprot -o '+output_file_org_dist+' --output-format=phylip '+dir_path+file_name)
		
			#Create trees from removed distance matrix
			output_file_remove_tree=path_remove_tree_sub_dir+file_name.rpartition('.')[0]+'.tree'
			os.system('/info/appbio10/bin/fnj --output-format=newick --input-format=phylip_dm -o '+output_file_remove_tree+' '+output_file_remove_dist)

			#Create trees from org distance matrix
			output_file_org_tree=path_org_tree_sub_dir+file_name.rpartition('.')[0]+'.tree'
			os.system('/info/appbio10/bin/fnj --output-format=newick --input-format=phylip_dm -o '+output_file_org_tree+' '+output_file_org_dist)

			#Read trees with dendropy
			tree_removed=dendropy.Tree.get_from_path(output_file_org_tree, schema='newick')
			tree_org=dendropy.Tree.get_from_path(output_file_remove_tree, schema='newick')
			tree_removed_list.append(tree_removed)
			tree_org_list.append(tree_org)

			#Insert trees into sql database
			[test,tree_id]=check_tree(tree_removed,sql_cursor) #Check if tree is already in database and get the id if it is or the maximum id+1 if it isnt
			if not test:
				t=(tree_id,str(tree_removed),)
				sql_cursor.execute('''insert into Tree values(?,?)''',t)
			tree_removed_id_list.append(tree_id)

			[test,tree_id]=check_tree(tree_org,sql_cursor) #Check if tree is already in database and get the id if it is or the maximum id+1 if it isnt
			if not test:
				t=(tree_id,str(tree_org),)
				sql_cursor.execute('''insert into Tree values(?,?)''',t)
			tree_org_id_list.append(tree_id)

			#Move multialignemnt with removed columns
			name=file_name.rpartition('.')[0] #Remove .msl from filename
			shutil.move(path_new_file,path_remove_align_sub_dir+name+'.removed.msl')

	if not tree_base_bool:
		sys.stderr.write('No base tree\n')
		sys.exit()

	#Get distances between trees
	for i in range(len(tree_removed_list)):
		t_removed=(expr_id,tree_base_id,tree_removed_id_list[i],tree_base.symmetric_difference(tree_removed_list[i]))
		t_org=(expr_id,tree_base_id,tree_org_id_list[i],tree_base.symmetric_difference(tree_org_list[i]))
		sql_cursor.execute('''insert into Distances values(?,?,?,?,'removed')''',t_removed) #Put into database
		sql_cursor.execute('''insert into Distances values(?,?,?,?,'original')''',t_org)

	#Calculate and store statistics
	#Get the distances from the database
	removed_distance_list=[]
	org_distance_list=[]
	t=(expr_id,)
	sql_cursor.execute('''select distance from Distances where type like 'removed' and expr_id=? ''',t)
	for row in sql_cursor:
		removed_distance_list.append(row[0])
	sql_cursor.execute('''select distance from Distances where type like 'original' and expr_id=? ''',t)
	for row in sql_cursor:
		org_distance_list.append(row[0])

	#Calculate statistics
	t=(expr_id,min(removed_distance_list),min(org_distance_list),max(removed_distance_list),max(org_distance_list),numpy.median(removed_distance_list),numpy.median(org_distance_list),numpy.mean(removed_distance_list),numpy.mean(org_distance_list))
	sql_cursor.execute('''insert into Statistics values(?,?,?,?,?,?,?,?,?)''',t)

sql_database.commit()
sql_cursor.close()
sys.stderr.write('Done\n')
